import numpy as np
import pandas as pd
from datetime import timedelta
import os
import joblib
from models import Data
from statsmodels.tsa.statespace.sarimax import SARIMAX

# 模型保存路径
MODEL_DIR = "models"
os.makedirs(MODEL_DIR, exist_ok=True)

# 加载200个最新数据
def load_data():
    data_records = Data.query.order_by(Data.timestamp.desc()).limit(200).all()
    df = pd.DataFrame([{
        'timestamp': record.timestamp,
        'temp': record.temp,
        'humidity': record.humidity,
        'light': record.light
    } for record in data_records])

    df.set_index('timestamp', inplace=True)
    df.sort_index(inplace=True)
    return df


# 训练模型
def train_sarima_model(df):
    order = (2, 1, 1)  # ARIMA 参数

    model_temp = SARIMAX(df['temp'], order=order)
    results_temp = model_temp.fit(disp=False)

    model_humidity = SARIMAX(df['humidity'], order=order)
    results_humidity = model_humidity.fit(disp=False)

    model_light = SARIMAX(df['light'], order=order)
    results_light = model_light.fit(disp=False)

    return results_temp, results_humidity, results_light


# 保存模型到本地
def save_models(results_temp, results_humidity, results_light):
    joblib.dump(results_temp, os.path.join(MODEL_DIR, 'temp_model.pkl'))
    joblib.dump(results_humidity, os.path.join(MODEL_DIR, 'humidity_model.pkl'))
    joblib.dump(results_light, os.path.join(MODEL_DIR, 'light_model.pkl'))


# 尝试加载模型，如果不存在则训练并保存
def get_or_train_models(df):
    temp_path = os.path.join(MODEL_DIR, 'temp_model.pkl')
    humidity_path = os.path.join(MODEL_DIR, 'humidity_model.pkl')
    light_path = os.path.join(MODEL_DIR, 'light_model.pkl')

    if os.path.exists(temp_path) and os.path.exists(humidity_path) and os.path.exists(light_path):
        # 加载已有模型
        results_temp = joblib.load(temp_path)
        results_humidity = joblib.load(humidity_path)
        results_light = joblib.load(light_path)
        return results_temp, results_humidity, results_light
    else:
        # 没有模型，重新训练并保存
        results_temp, results_humidity, results_light = train_sarima_model(df)
        save_models(results_temp, results_humidity, results_light)
        return results_temp, results_humidity, results_light


# 预测未来值
def predict_future(results_temp, results_humidity, results_light, steps=30):
    df = load_data()

    # 获取最后一条时间记录
    last_time = df.index[-1]

    # 生成未来 30 个时间点（每分钟一个）
    future_timestamps = [last_time + timedelta(minutes=i) for i in range(1, steps + 1)]

    # 获取预测值
    pred_temp = results_temp.get_forecast(steps=steps)
    pred_humidity = results_humidity.get_forecast(steps=steps)
    pred_light = results_light.get_forecast(steps=steps)

    predicted_temps = pred_temp.predicted_mean.tolist()
    predicted_humidities = pred_humidity.predicted_mean.tolist()
    predicted_light = pred_light.predicted_mean.tolist()

    # 返回预测值和对应时间戳
    predictions = list(zip(predicted_temps, predicted_humidities, predicted_light))
    return future_timestamps, predictions